#!/usr/bin/python
# -*- coding:utf-8 -*-
import redis
import logging

log = logging.getLogger(__name__)


class ProxyManager(object):
    def __init__(self):
        self.redis = redis.Redis()

    def has_proxy(self):
        """
        检查是否存在代理ip
        :return: True/False
        """
        return self.redis.zcard('proxy:new') != 0

    def remove_proxies(self, proxies):
        return self.redis.zrem('proxy:new', *proxies)

    def add_proxy(self, proxy):
        if not proxy:
            return
        self.redis.zadd('proxy:new', proxy, 1)

    def add_proxies(self, proxies):
        if proxies is None or len(proxies) == 0:
            return
        for proxy in proxies:
            self.add_proxy(proxy)

    def get_proxy(self):
        if not self.has_proxy():
            return None
        try:
            proxy = self.redis.zrevrange('proxy:new', 0, 0).pop()
            self.redis.zincrby('proxy:new', proxy, -1)
            return proxy.decode('utf-8')
        except IndexError as e:
            log.error(e)
            return None

    def get_proxies(self, num):
        """
        返回一个迭代器
        :param num:
        :return:
        """
        for i in range(num):
            if not self.has_proxy():
                return
            yield self.get_proxy()

    def get_all(self):
        return list(map(lambda p: p.decode('utf-8'), self.redis.zrevrange('proxy:new', 0, -1)))


manager = ProxyManager()
